package words;

import java.util.*;
import java.io.*;


public class IntexSentence {

  // sgml IDs
  static public String STR_P1 = "p1";
  static public String STR_P2 = "p2";
  static public String STR_PAIR = "pair";

  // sgml sentence
  SgmlTokSentence m_sgmlTS;
  // array of sentence entities
  Annotation[] m_entities;
  // upper-triangular matrix of entity pairs
  EntityPair[][] m_pairs;


  public IntexSentence(SgmlTokSentence sgmlTS, String strTag)
  {
    m_sgmlTS = sgmlTS;
    
    int nIdTag = m_sgmlTS.m_tagDef.getId(strTag);
    int nIdP1 = m_sgmlTS.m_tagDef.getId(STR_P1);
    int nIdP2 = m_sgmlTS.m_tagDef.getId(STR_P2);

    // isolate entities
    TreeSet setTags = new TreeSet();
    setTags.add(new Integer(nIdTag));
    setTags.add(new Integer(nIdP1));
    setTags.add(new Integer(nIdP2));

    LinkedList lstEntities = sgmlTS.getAnnotations(setTags);
    m_entities = new Annotation[lstEntities.size()];
    lstEntities.toArray(m_entities);

    // create matrix of entity pairs
    m_pairs = new EntityPair[m_entities.length][m_entities.length];

    // populate matrix with positive pairs
    ListIterator it1 = sgmlTS.m_lstAnno.listIterator();
    while (it1.hasNext()) {
      Annotation anno1 = (Annotation) it1.next();
      if (anno1.m_nId == nIdP1) {
	String strPair1 = anno1.getAttribute(STR_PAIR);
	ListIterator it2 = sgmlTS.m_lstAnno.listIterator(it1.nextIndex());
	while (it2.hasNext()) {
	  Annotation anno2 = (Annotation) it2.next();
	  // ignore overlapping entities
	  if (!anno2.includes(anno1) && !anno1.includes(anno2)) {
	    if (anno2.m_nId == nIdP2) {
	      String strPair2 = anno2.getAttribute(STR_PAIR);
	      if (strPair1.equals(strPair2)) {
		int nPos1 = findEntity(anno1);
		int nPos2 = findEntity(anno2);
		m_pairs[nPos1][nPos2] = new EntityPair(m_sgmlTS.m_nSent, 
						       nPos1, nPos2);
		m_pairs[nPos1][nPos2].setTrueValue(true);
	      }
	    }
	  }
	}
      }
    }

    // complete matrix with negative entity pairs
    for (int i = 0; i < m_entities.length; i++) {
      Annotation anno1 = m_entities[i];
      for (int j = i + 1; j < m_entities.length; j++) {
	Annotation anno2 = m_entities[j];
	// ignore overlapping entities
	if (!anno2.includes(anno1) && !anno1.includes(anno2)) {
	  if (m_pairs[i][j] == null) {
	    m_pairs[i][j] = new EntityPair(m_sgmlTS.m_nSent, i, j);
	    m_pairs[i][j].setTrueValue(false);
	  }
	}
      }
    }
  }


  protected int findEntity(Annotation anno)
  {
    // assume entities are ordered based on m_nBegin
    for (int i = 0; i < m_entities.length; i++) {
      //      if (m_entities[i].m_nBegin > anno.m_nBegin)
      //	return -1;
      if (m_entities[i].m_nBegin == anno.m_nBegin &&
	  m_entities[i].m_nLen == anno.m_nLen)
	return i;
    }
    
    return -1;
  }


  public EntityPair getEntityPair(int i, int j)
  {
    int p1 = i;
    int p2 = j;
    if (i > j) {
      p1 = j;
      p2 = i;
    }

    return m_pairs[p1][p2];
  }


  public int entCount()
  {
    return m_entities.length;
  }


  public int[] pairsCounts()
  {
    int[] counts = new int[2];
    for (int i = 0; i < m_entities.length; i++)
      for (int j = i + 1; j < m_entities.length; j++)
	if (m_pairs[i][j] != null)
	  counts[m_pairs[i][j].m_nTrueLabel]++;
	    
    return counts;
  }

  public Annotation[] getEntities()
  {
    return m_entities;
  }

  public Annotation getEntity(int nPos)
  {
    return m_entities[nPos];
  }


  public Annotation getClosestLEntity(Annotation anno)
  {
    Annotation annoMax = null;
    int endMax = 0;
    for (int i = 0; i < m_entities.length; i++) {
      int end = m_entities[i].m_nBegin + m_entities[i].m_nLen;
      if (end <= anno.m_nBegin && end > endMax) {
	endMax = end;
	annoMax = m_entities[i];
      }
    }

    return annoMax;
  }


  public int writeInf(PrintWriter pw, int crtPairs)
  {
    int nPairs = 0;
    
    LinkedList lstAnno = new LinkedList();
    // add entities
    for (int i = 0; i < m_entities.length; i++)
      lstAnno.add(m_entities[i]);
    
    // add inferred pairs
    for (int i = 0; i < m_entities.length; i++)
      for (int j = i + 1; j < m_entities.length; j++) {
	if (m_pairs[i][j] != null/* && m_pairs[i][j].isInfPos()*/) {
	  double conf = 0.0;
	  if (m_pairs[i][j].isInfPos())
	    conf = m_pairs[i][j].m_conf;
	  else
	    conf = 1 - m_pairs[i][j].m_conf;
	  if (conf > 0.1) {
	    // add annotation for first element of the pair
	    Annotation anno1 = new Annotation(m_sgmlTS.m_tagDef, STR_P1, 
					      m_entities[i].m_nBegin, 
					      m_entities[i].m_nLen);
	    anno1.setAttribute(STR_PAIR, Integer.toString(crtPairs + nPairs));
	    anno1.setAttribute("conf", Double.toString(conf));
	    lstAnno.add(anno1);
	    // add annotation for second element of the pair
	    Annotation anno2 = new Annotation(m_sgmlTS.m_tagDef, STR_P2, 
					      m_entities[j].m_nBegin, 
					      m_entities[j].m_nLen);
	    anno2.setAttribute(STR_PAIR, Integer.toString(crtPairs + nPairs));
	    anno2.setAttribute("conf", Double.toString(conf));
	    lstAnno.add(anno2);
	    // increment number of pairs
	    nPairs++;
	  }
	}
      }
    
    // write annotations to stream
    m_sgmlTS.setAnnotations(lstAnno);
    pw.println(m_sgmlTS.toString());
    
    return nPairs;
  }
  

  public void writeACE(StringBuffer strText,
		       StringBuffer strEnts,
		       StringBuffer strRels,
		       int[] counts, int nHeader)
  {
    int nText = strText.length() + nHeader;
    int nEnts = counts[0];
    int nRels = counts[1];

    // text document position for each sentence token
    int[] m_index = new int[m_sgmlTS.m_vecTokens.length];

    // write raw text
    String strSent = new String();
    int crtPos = nText;
    for (int i = 0; i < m_sgmlTS.m_vecTokens.length; i++) {
      strSent += m_sgmlTS.m_vecTokens[i] + " ";
      m_index[i] = crtPos;
      crtPos += m_sgmlTS.m_vecTokens[i].length() + 1;
    }
    strSent += "\n";
    strText.append(strSent);

    // write entities
    for (int i = 0; i < m_entities.length; i++) {
      int nFirst = m_entities[i].m_nBegin;
      int nLast = m_entities[i].m_nBegin + m_entities[i].m_nLen - 1;
      int nStart = m_index[nFirst];
      int nEnd = m_index[nLast] + m_sgmlTS.m_vecTokens[nLast].length() - 1;
      String strEntID = "E" + String.valueOf(nEnts + i);
      String strEntText = strSent.substring(nStart - nText, nEnd - nText + 1);
      strEnts.append(getEntityText(strEntID, strEntText, nStart, nEnd)); 
    }

    // write relations
    for (int i = 0; i < m_entities.length; i++) {
      Annotation anno1 = m_entities[i];
      String strEntID1 = "E" + String.valueOf(nEnts + i);
      for (int j = i + 1; j < m_entities.length; j++) {
	Annotation anno2 = m_entities[j];
	String strEntID2 = "E" + String.valueOf(nEnts + j);
	if (m_pairs[i][j] != null && m_pairs[i][j].isTruePos()) {
	  String strRelID = "R" + String.valueOf(nRels);
	  nRels++;
	  strRels.append(getRelationText(strRelID, strEntID1, strEntID2));
	}
      }
    }

    counts[0] = nEnts + m_entities.length;
    counts[1] = nRels;
  }


  protected String getEntityText(String strEntID, String strEntText, 
				 int nStart, int nEnd)
  {
    String strText = new String();
    strText += "    <entity ID=\"" + strEntID + "\">\n";
    strText += "      <entity_type GENERIC=\"FALSE\">PERSON</entity_type>\n";
    strText += "        <entity_mention TYPE=\"NAME\" ID=\"" + strEntID + "-M\">\n";
    strText += "          <extent>\n";
    strText += "            <charseq>\n";
    strText += "            <!-- string = \"" + strEntText + "\" -->\n";
    strText += "               <start>" + String.valueOf(nStart) + "</start><end>" + String.valueOf(nEnd) + "</end></charseq>\n";
    strText += "          </extent>\n";
    strText += "          <head>\n";
    strText += "            <charseq>\n";
    strText += "            <!-- string = \"" + strEntText + "\" -->\n";
    strText += "               <start>" + String.valueOf(nStart) + "</start><end>" + String.valueOf(nEnd) + "</end></charseq>\n";
    strText += "          </head>\n";
    strText += "        </entity_mention>\n";
    strText += "      </entity>\n\n\n";

    return strText;
  }


  protected String getRelationText(String strRelID, 
				   String strEntID1, String strEntID2)
  {
    String strText = new String();
    strText += "    <relation ID=\"" + strRelID + "\" TYPE=\"SOC\" SUBTYPE=\"SIBLING\" CLASS=\"EXPLICIT\">\n";
    strText += "      <rel_entity_arg ENTITYID=\"" + strEntID1 + "\" ARGNUM=\"1\">\n";
    strText += "        </rel_entity_arg>\n";
    strText += "      <rel_entity_arg ENTITYID=\"" + strEntID2 + "\" ARGNUM=\"2\">\n";
    strText += "        </rel_entity_arg>\n";
    strText += "      <relation_mentions>\n";
    strText += "        <relation_mention ID=\"" + strRelID + "-M\">\n";
    strText += "          <rel_mention_arg MENTIONID=\"" + strEntID1 + "-M\" ARGNUM=\"1\">\n";
    strText += "            </rel_mention_arg>\n";
    strText += "          <rel_mention_arg MENTIONID=\"" + strEntID2 + "-M\" ARGNUM=\"2\">\n";
    strText += "            </rel_mention_arg>\n";
    strText += "          </relation_mention>\n";
    strText += "      </relation_mentions>\n";
    strText += "    </relation>\n\n";

    return strText;
  }


}
